import socket
import threading
import cv2
import numpy as np
import mediapipe as mp
from threading import Lock
import json

# ================= 可调参数区域 =================
# 服务器配置
HOST = "127.0.0.1"
PORT = 12345
PORT_C = 12346

from threading import Lock

# 全局参数存储（带线程锁保护）
GLOBAL_PARAMS = {
    "POSE_PARAMS": {
        "min_detection_confidence": 0.5,
        "min_tracking_confidence": 0.5,
        "model_complexity": 1
    },
    "LANDMARK_SPEC": {
        "color": (0, 255, 0),
        "thickness": 2,
        "circle_radius": 2
    },
    "CONNECTION_SPEC": {
        "color": (0, 5, 0),
        "thickness": 2
    },
    "IMAGE_PARAMS": {
        "image_quality": 90
        
    },
}

params_lock = Lock()  # 参数访问互斥锁

def validate_parameters(new_params):
    """参数有效性验证[8](@ref)"""
    # 模型复杂度校验
    if 'model_complexity' in new_params.get('POSE_PARAMS', {}):
        if not 0 <= new_params['POSE_PARAMS']['model_complexity'] <= 2:
            raise ValueError("模型复杂度需在0-2范围内")
    
    # 置信度阈值校验
    for conf_key in ['min_detection_confidence', 'min_tracking_confidence']:
        if conf_key in new_params.get('POSE_PARAMS', {}):
            value = new_params['POSE_PARAMS'][conf_key]
            if not 0 <= value <= 1:
                raise ValueError(f"{conf_key}需在0-1范围内")
    
    # JPEG质量校验
    if 'IMAGE_PARAMS' in new_params:
        image_params = new_params['IMAGE_PARAMS']
        if 'image_quality' in image_params:
            iq = image_params['image_quality']
            if not 0 <= iq <= 100:
                raise ValueError("图像质量需在0-100范围内")
    
    return True
# ================= 参数区域结束 =================

# ================= 参数控制服务端 =================

def param_control_server():
    """参数控制服务线程[1,2](@ref)"""
    param_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    param_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    param_socket.bind((HOST, PORT_C))  # 使用相邻端口
    param_socket.listen(5)
    
    print(f"参数控制服务已启动 ({HOST}:{PORT_C})...")
    
    while True:
        conn, addr = param_socket.accept()
        #print("等待参数发送")
        try:
            # 接收完整JSON数据
            raw_data = conn.recv(4096).decode('utf-8')
            new_params = json.loads(raw_data)
            print("等待参数发送")
            if validate_parameters(new_params):
                with params_lock:
                    # 分层更新参数[8](@ref)
                    for category, values in new_params.items():
                        if category in GLOBAL_PARAMS:
                            GLOBAL_PARAMS[category].update(values)
                    print("参数更新成功:", GLOBAL_PARAMS)
                conn.sendall(b"OK")
            else:
                conn.sendall(b"Invalid Parameters")
                
        except Exception as e:
            print(f"参数更新错误: {str(e)}")
            conn.sendall(f"ERROR: {str(e)}".encode())
        finally:
            conn.close()
# ================= 参数控制服务端结束 =================

# MediaPipe 初始化
mp_drawing = mp.solutions.drawing_utils
mp_pose = mp.solutions.pose

def process_image(image_data):
    """改造后的图像处理函数"""
    # 获取当前参数快照[5](@ref)
    with params_lock:
        pose_params = GLOBAL_PARAMS["POSE_PARAMS"].copy()
        landmark_spec = GLOBAL_PARAMS["LANDMARK_SPEC"].copy()
        connection_spec = GLOBAL_PARAMS["CONNECTION_SPEC"].copy()
        jpeg_quality = GLOBAL_PARAMS["IMAGE_PARAMS"]["image_quality"]  # 修改这里
    
    # 处理流程（保持原有逻辑）
    img_array = np.frombuffer(image_data, dtype=np.uint8)
    img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
    
    with mp_pose.Pose(**pose_params) as pose:
        image = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        results = pose.process(image)
        
        if results.pose_landmarks:
            mp_drawing.draw_landmarks(
                image, 
                results.pose_landmarks,
                mp_pose.POSE_CONNECTIONS,
                landmark_drawing_spec=mp_drawing.DrawingSpec(**landmark_spec),
                connection_drawing_spec=mp_drawing.DrawingSpec(**connection_spec))
        
        processed_image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
        _, jpeg = cv2.imencode('.jpg', processed_image, 
                             [int(cv2.IMWRITE_JPEG_QUALITY), jpeg_quality])
        return jpeg.tobytes()
    


def handle_client(conn):
    try:
        # 接收所有数据
        data = b''
        while True:
            chunk = conn.recv(4096)
            if not chunk:
                break
            data += chunk
        
        # 处理图像
        processed_data = process_image(data)
        
        # 返回处理结果
        conn.sendall(processed_data)
    except Exception as e:
        print(f"处理错误: {str(e)}")
    finally:
        conn.close()

def start_server():
    # 启动参数控制线程[7](@ref)
    threading.Thread(target=param_control_server, daemon=True).start()
    
    # 原有图像服务逻辑
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        s.bind((HOST, PORT))
        s.listen(5)
        print(f"Python骨骼检测服务已启动 ({HOST}:{PORT})...")
        
        while True:
            conn, addr = s.accept()
            print(f"新连接: {addr}")
            threading.Thread(target=handle_client, args=(conn,)).start()



if __name__ == "__main__":
    start_server()